package cc.git.liuyan.customeraiagent.core.suppertool;

import cc.git.liuyan.customeraiagent.core.embeddingmodel.EmbeddingModelInputData;
import cc.git.liuyan.customeraiagent.core.embeddingmodel.EmbeddingModelOutputData;
import cc.git.liuyan.customeraiagent.core.embeddingmodel.HybridEmbeddingModel;
import cc.git.liuyan.customeraiagent.core.vectordatabase.KnowledgeVectorData;
import cc.git.liuyan.customeraiagent.core.vectordatabase.VectorDataBase;
import lombok.extern.slf4j.Slf4j;

import java.math.BigDecimal;
import java.util.List;

//知识库超级工具
@Slf4j
public class KnowledgeSupperTool {
    private static HybridEmbeddingModel innerHybridEmbeddingModel;
    private static VectorDataBase innerVectorDataBase;

    public static void init(HybridEmbeddingModel hybridEmbeddingModel, VectorDataBase vectorDataBase) {
        innerHybridEmbeddingModel = hybridEmbeddingModel;
        innerVectorDataBase = vectorDataBase;
    }

    /**
     * 训练知识库，并返回分片总数
     */
    public static Long train(Long knowledgeId, EmbeddingModelInputData inputData) {
        List<EmbeddingModelOutputData> allSegmentVectors = innerHybridEmbeddingModel.embed(inputData);
        log.info("KnowledgeSupperTool.train 分片总数{}", allSegmentVectors.size());
        innerVectorDataBase.createKnowledgeVectorCollection(innerHybridEmbeddingModel.getDenseEmbeddingModel().dimension(), innerHybridEmbeddingModel.getSpaseEmbeddingModel().dimension());
        innerVectorDataBase.deleteKnowledgeTrain(knowledgeId);
        Long segmentNum = 1L;
        for (EmbeddingModelOutputData segmentVector : allSegmentVectors) {
            KnowledgeVectorData knowledgeVectorData = new KnowledgeVectorData();
            knowledgeVectorData.setKnowledgeId(knowledgeId);
            knowledgeVectorData.setSegmentNum(segmentNum);
            knowledgeVectorData.setSegmentContent(segmentVector.getSegmentContent());
            knowledgeVectorData.setDenseVectors(segmentVector.getDenseVectors());
            knowledgeVectorData.setSparseVectors(segmentVector.getSparseVectors());
            knowledgeVectorData.setScore(1L);
            innerVectorDataBase.trainKowledge(knowledgeVectorData);
            log.info("KnowledgeSupperTool.train 分片:{}", segmentNum);
            segmentNum++;
        }
        log.info("KnowledgeSupperTool.train end success");
        return segmentNum;
    }

    public static EmbeddingModelOutputData question2Vectors(String userQuestion) {
        return innerHybridEmbeddingModel.embed(new EmbeddingModelInputData("remark", userQuestion)).get(0);
    }

    public static List<String> relationKnowledgeSearch(List<Long> knowledgeIds, List<BigDecimal> denseVectors, List<BigDecimal> sparseVectors) {
        return innerVectorDataBase.relationKnowledgeSearch(knowledgeIds, denseVectors, sparseVectors);
    }
}
